package com.tlgen.orm.utils.orm;

import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.tlgen.orm.param.MiddleParams;
import com.tlgen.orm.param.RelationParams;
import com.tlgen.orm.utils.log.SQLLogger;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.core.BeanPropertyRowMapper;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;

import static com.tlgen.orm.enums.Anno.ASSOCIATION;

// 优化后的关联关系处理器
public class AssociationProcessor {

    private static final Map<Class<?>, String> TABLE_NAME_CACHE = new ConcurrentHashMap<>();
    private static final Map<Class<?>, List<RelationParams>> ASSOCIATION_CACHE = new ConcurrentHashMap<>();
    private static final Map<Class<?>, List<RelationParams>> COLLECTION_CACHE = new ConcurrentHashMap<>();
    private static final Map<Class<?>, List<MiddleParams>> MIDDLE_MAPPING_CACHE = new ConcurrentHashMap<>();

    private static final Logger log = LoggerFactory.getLogger(AssociationProcessor.class);
    private static final String SELECT_PREFIX = "SELECT * FROM ";

    // 处理一对一关联关系（安全+高性能）
    public static <T> void processAssociations(List<JSONObject> jsonObjects, String primaryKey,
                                                Class<T> tClass, JdbcTemplate jdbcTemplate) {
        List<RelationParams> aList = ASSOCIATION_CACHE.computeIfAbsent(tClass, ORMUtils::getAssociationFields);
        if (CollectionUtils.isEmpty(aList)) return;

        // 并行流处理所有关联关系
        aList.parallelStream().forEach(relation -> {
            if (!ASSOCIATION.getValue().equals(relation.getRelationType())) return;

            try {
                Class<?> relationClass = Class.forName(relation.getColumnType());

                // 1. 批量收集外键值
                Set<Object> foreignKeyValues = jsonObjects.stream()
                        .map(json -> json.get(primaryKey))
                        .filter(Objects::nonNull)
                        .collect(Collectors.toSet());

                if (foreignKeyValues.isEmpty()) return;

                // 2. 智能解析外键字段名称
                String foreignKeyField = resolveForeignKey(relation, tClass, relationClass);

                // 3. 批量查询关联表数据
                String relationTable = TABLE_NAME_CACHE.computeIfAbsent(relationClass, ORMUtils::getTableName);
                String sql = buildBatchQuery(relationTable, foreignKeyField, foreignKeyValues);

                List<?> relationList = executeBatchQuery(jdbcTemplate, sql, foreignKeyValues, relationClass);

                // 4. 构建关联对象映射
                Map<Object, Object> relationMap = buildRelationMap(relationList, foreignKeyField);

                // 5. 填充关联对象
                fillAssociations(jsonObjects, primaryKey, relation.getColumnName(), relationMap);

            } catch (ClassNotFoundException e) {
                log.error("Association class not found: " + relation.getColumnType(), e);
            } catch (Exception e) {
                log.error("Error processing association: " + relation.getColumnName(), e);
            }
        });
    }

    // 智能解析外键字段名
    private static String resolveForeignKey(RelationParams relation, Class<?> masterClass, Class<?> relationClass) {
        // 优先使用注解指定的关联字段
        if (StringUtils.isNotBlank(relation.getRelationName())) {
            return ORMUtils.getColumnName(relationClass, relation.getRelationName());
        }

        // 默认规则：主表类名小写 + "Id"
        String defaultKey = masterClass.getSimpleName().toLowerCase() + "Id";
        String column = ORMUtils.getColumnName(relationClass, defaultKey);

        // 验证字段存在性
        if (column != null) {
            return column;
        }

        // 尝试直接使用主键名
        String primaryKey = ORMUtils.getPrimaryKey(masterClass);
        column = ORMUtils.getColumnName(relationClass, primaryKey);
        if (column != null) {
            return column;
        }

        throw new RuntimeException("No valid foreign key found for association between "
                + masterClass.getSimpleName() + " and " + relationClass.getSimpleName() +
                ". Tried: " + defaultKey + ", " + primaryKey);
    }

    // 构建批量查询SQL
    private static String buildBatchQuery(String table, String field, Set<Object> values) {
        String placeholders = values.stream()
                .map(v -> "?")
                .collect(Collectors.joining(","));

        return SELECT_PREFIX + table + " WHERE " + field + " IN (" + placeholders + ")";
    }

    // 构建关联对象映射
    private static Map<Object, Object> buildRelationMap(List<?> relationList, String foreignKeyColumn) {
        Map<Object, Object> map = new ConcurrentHashMap<>();

        for (Object obj : relationList) {
            // 使用 ORMUtils.getFieldValue 方法，它会自动处理列名到字段名的转换
            Object keyValue = ORMUtils.getFieldValue(obj, foreignKeyColumn);
            if (keyValue != null) {
                // 处理重复记录：一对一只取第一条
                String key = String.valueOf(keyValue);
                if (map.containsKey(key)) {
                    log.warn("One-to-One association found multiple records for key: {} in {}. Using first record found.",
                            keyValue, obj.getClass().getSimpleName());
                } else {
                    map.put(key, obj);
                }
            }
        }
        return map;
    }

    // 填充关联对象
    private static void fillAssociations(List<JSONObject> jsonObjects, String primaryKey,
                                         String relationField, Map<Object, Object> relationMap) {
        for (JSONObject json : jsonObjects) {
            Object masterId = json.get(primaryKey);
            String key = String.valueOf(masterId);
            Object relationObj = relationMap.get(key);
            if (relationObj != null) {
                json.put(relationField, relationObj);
            }
        }
    }


    // 处理一对多关联关系（高性能批量查询）
    public static <T> void processCollections(List<JSONObject> jsonObjects, String primaryKey,
                                               Class<T> tClass, JdbcTemplate jdbcTemplate) {
        List<RelationParams> collectionList = COLLECTION_CACHE.computeIfAbsent(tClass, ORMUtils::getCollectionFields);
        if (CollectionUtils.isEmpty(collectionList)) return;

        // 1. 收集所有主表ID
        Set<Object> masterIds = jsonObjects.stream()
                .map(json -> json.get(primaryKey))
                .filter(Objects::nonNull)
                .collect(Collectors.toSet());

        if (masterIds.isEmpty()) return;

        // 2. 并行处理所有集合关系
        collectionList.parallelStream().forEach(collection -> {
            try {
                // 判断处理模式：有中间表 or 无中间表
                if (collection.getHasMiddleClass()) {
                    // 处理有中间表的情况
                    processMiddleTableCollection(collection, tClass, masterIds, jsonObjects, primaryKey, jdbcTemplate);
                } else {
                    // 处理无中间表的情况
                    processDirectForeignKeyCollection(collection, tClass, masterIds, jsonObjects, primaryKey, jdbcTemplate);
                }
            } catch (Exception e) {
                log.error("Error processing collection: " + collection.getColumnName(), e);
            }
        });
    }

    /**
     * 处理有中间表的一对多关联
     */
    private static void processMiddleTableCollection(RelationParams collection, Class<?> masterClass,
                                                     Set<Object> masterIds, List<JSONObject> jsonObjects,
                                                     String primaryKey, JdbcTemplate jdbcTemplate) {
        // 1. 解析中间表关系
        MiddleMapping mapping = resolveMiddleMapping(collection, masterClass);
        if (mapping == null) return;

        // 2. 批量查询中间表映射
        Map<Object, Set<Object>> masterToTargetMap = queryMiddleTable(
                jdbcTemplate, mapping, masterIds
        );

        // 3. 批量查询目标对象
        Set<Object> targetIds = extractDistinctTargetIds(masterToTargetMap);
        if (targetIds.isEmpty()) return;

        List<?> targetList = queryTargetObjects(
                jdbcTemplate, mapping.targetClass, targetIds
        );

        // 4. 构建目标对象映射
        Map<Object, Object> targetMap = buildTargetMap(
                targetList, mapping.targetClass
        );

        // 5. 填充集合对象
        fillCollections(jsonObjects, primaryKey,
                collection.getColumnName(), masterToTargetMap, targetMap);
    }

    /**
     * 处理无中间表（直接外键）的一对多关联
     */
    private static void processDirectForeignKeyCollection(RelationParams collection, Class<?> masterClass,
                                                          Set<Object> masterIds, List<JSONObject> jsonObjects,
                                                          String primaryKey, JdbcTemplate jdbcTemplate) {
        try {
            // 1. 解析目标类型
            Class<?> targetClass = Class.forName(collection.getColumnType());
            String targetTable = TABLE_NAME_CACHE.computeIfAbsent(targetClass, ORMUtils::getTableName);

            // 2. 解析外键字段名
            String foreignKey = resolveForeignKeyForCollection(collection, masterClass, targetClass);

            // 3. 批量查询目标对象
            String sql = SELECT_PREFIX + targetTable + " WHERE " + foreignKey + " IN (" +
                    masterIds.stream().map(id -> "?").collect(Collectors.joining(",")) + ")";

            List<?> targetList = executeBatchQuery(jdbcTemplate, sql, masterIds, targetClass);

            // 4. 按外键值分组目标对象
            Map<Object, List<Object>> masterToTargetsMap = new ConcurrentHashMap<>();

            targetList.parallelStream().forEach(targetObj -> {
                try {
                    Object fkValue = ORMUtils.getFieldValue(targetObj, foreignKey);
                    if (fkValue != null) {
                        String key = String.valueOf(fkValue);
                        masterToTargetsMap.computeIfAbsent(key, k -> new ArrayList<>())
                                .add(targetObj);
                    }
                } catch (Exception e) {
                    log.error("Failed to get foreign key value from {}: {}", targetClass.getSimpleName(), e.getMessage());
                }
            });

            // 5. 填充集合对象
            jsonObjects.forEach(json -> {
                Object masterId = json.get(primaryKey);
                String key = String.valueOf(masterId);
                if (masterId != null && masterToTargetsMap.containsKey(key)) {
                    json.put(collection.getColumnName(), masterToTargetsMap.get(key));
                }
            });

        } catch (ClassNotFoundException e) {
            log.error("Target class not found: " + collection.getColumnType(), e);
        }
    }

    /**
     * 中间表映射关系解析
     */
    private static MiddleMapping resolveMiddleMapping(RelationParams collection, Class<?> masterClass) {
        List<MiddleParams> mList = MIDDLE_MAPPING_CACHE.computeIfAbsent(
                collection.getRelationClass(), ORMUtils::getMiddleMappingParams);

        if (CollectionUtils.isEmpty(mList)) return null;

        // 区分源表和目标表
        MiddleParams source = null, target = null;
        for (MiddleParams p : mList) {
            if (p.getMappingClass().equals(masterClass)) {
                source = p;
            } else {
                target = p;
            }
        }

        if (source == null || target == null) {
            log.warn("Invalid middle mapping for collection: {}", collection.getColumnName());
            return null;
        }

        return new MiddleMapping(
                source.getTableName(),
                ORMUtils.getColumnName(source.getMappingClass(), source.getMappingColumn()),
                target.getMappingClass(),
                ORMUtils.getColumnName(target.getMappingClass(), target.getMappingColumn())
        );
    }

    /**
     * 查询中间表关系
     */
    private static Map<Object, Set<Object>> queryMiddleTable(JdbcTemplate jdbcTemplate,
                                                             MiddleMapping mapping,
                                                             Set<Object> masterIds) {
        String sql = String.format(
                "SELECT %s, %s FROM %s WHERE %s IN (%s)",
                mapping.sourceColumn,
                mapping.targetColumn,
                mapping.tableName,
                mapping.sourceColumn,
                masterIds.stream().map(k -> "?").collect(Collectors.joining(","))
        );

        SQLLogger.printSQL(sql, Lists.newArrayList(masterIds));

        List<Map<String, Object>> rows = jdbcTemplate.queryForList(sql, masterIds.toArray());

        // 按源ID分组的目标ID集合
        Map<Object, Set<Object>> mappingMap = new ConcurrentHashMap<>();
        for (Map<String, Object> row : rows) {
            Object sourceId = row.get(mapping.sourceColumn);
            Object targetId = row.get(mapping.targetColumn);

            if (sourceId != null && targetId != null) {
                mappingMap.computeIfAbsent(sourceId, k -> ConcurrentHashMap.newKeySet())
                        .add(targetId);
            }
        }
        return mappingMap;
    }

    /**
     * 解析集合关系中的外键字段名（无中间表的情况）
     */
    private static String resolveForeignKeyForCollection(RelationParams collection, Class<?> masterClass,
                                                         Class<?> targetClass) {
        // 1. 如果注解中指定了relate属性，则使用该值
        if (StringUtils.isNotBlank(collection.getRelationName())) {
            return ORMUtils.getColumnName(targetClass, collection.getRelationName());
        }

        // 2. 默认规则：主表类名小写 + "Id"（如authorId）
        String defaultKey = masterClass.getSimpleName().toLowerCase() + "Id";
        String column = ORMUtils.getColumnName(targetClass, defaultKey);

        if (column != null) {
            return column;
        }

        // 3. 尝试使用主键字段名
        String primaryKey = ORMUtils.getPrimaryKey(masterClass);
        column = ORMUtils.getColumnName(targetClass, primaryKey);
        if (column != null) {
            return column;
        }

        // 4. 抛出异常，提示配置错误
        throw new RuntimeException("No valid foreign key found for direct one-to-many association between "
                + masterClass.getSimpleName() + " and " + targetClass.getSimpleName() +
                ". Tried: " + defaultKey + ", " + primaryKey +
                ". Please specify 'relate' attribute in @Collection annotation.");
    }

    /**
     * 中间表映射关系包装类
     */
    private static class MiddleMapping {
        final String tableName;
        final String sourceColumn;
        final Class<?> targetClass;
        final String targetColumn;

        MiddleMapping(String tableName, String sourceColumn,
                      Class<?> targetClass, String targetColumn) {
            this.tableName = tableName;
            this.sourceColumn = sourceColumn;
            this.targetClass = targetClass;
            this.targetColumn = targetColumn;
        }
    }

    // 通用的批量查询方法
    private static List<?> executeBatchQuery(JdbcTemplate jdbcTemplate, String sql,
                                             Set<Object> params, Class<?> targetClass) {
        SQLLogger.printSQL(sql, Lists.newArrayList(params));
        return jdbcTemplate.query(
                sql,
                new BeanPropertyRowMapper<>(targetClass),
                params.toArray()
        );
    }

    private static Set<Object> extractDistinctTargetIds(Map<Object, Set<Object>> mappingMap) {
        return mappingMap.values().stream()
                .flatMap(Set::stream)
                .collect(Collectors.toSet());
    }

    // -------------------------------- 缺失方法2：批量查询目标对象 ---------------------------------
    private static List<?> queryTargetObjects(JdbcTemplate jdbcTemplate,
                                              Class<?> targetClass,
                                              Set<Object> targetIds) {
        String table = TABLE_NAME_CACHE.computeIfAbsent(targetClass, ORMUtils::getTableName);
        String pk = ORMUtils.getPrimaryKey(targetClass);
        String pkColumn = ORMUtils.getColumnName(targetClass, pk);

        String sql = SELECT_PREFIX + table + " WHERE " + pkColumn + " IN (" +
                targetIds.stream().map(id -> "?").collect(Collectors.joining(",")) + ")";

        // 打印SQL日志（确保SQLLogger可用）
        SQLLogger.printSQL(sql, Lists.newArrayList(targetIds));

        return jdbcTemplate.query(
                sql,
                new BeanPropertyRowMapper<>(targetClass),
                targetIds.toArray()
        );
    }

    // -------------------------------- 缺失方法3：填充集合对象 ---------------------------------
    private static void fillCollections(List<JSONObject> jsonObjects,
                                        String primaryKey,
                                        String collectionField,
                                        Map<Object, Set<Object>> masterToTargetMap,
                                        Map<Object, Object> targetMap) {
        for (JSONObject json : jsonObjects) {
            Object masterId = json.get(primaryKey);
            if (masterId != null && masterToTargetMap.containsKey(masterId)) {
                Set<Object> targetIds = masterToTargetMap.get(masterId);
                if (targetIds != null && !targetIds.isEmpty()) {
                    List<Object> related = targetIds.stream()
                            .map(targetMap::get)
                            .filter(Objects::nonNull)
                            .collect(Collectors.toList());
                    json.put(collectionField, related);
                }
            }
        }
    }

    // -------------------------------- 以下是必需的基础方法 ---------------------------------

    // 构建目标对象映射
    private static Map<Object, Object> buildTargetMap(List<?> targetList, Class<?> targetClass) {
        String pk = ORMUtils.getPrimaryKey(targetClass);
        return targetList.parallelStream()
                .collect(Collectors.toConcurrentMap(
                        obj -> ORMUtils.getFieldValue(obj, pk),
                        Function.identity(),
                        (a, b) -> a  // 重复时保留第一个
                ));
    }

}